import tushare as ts
import pandas as pd
from datetime import datetime, timedelta
from typing import Optional, Dict
import logging
from django.conf import settings
from functools import lru_cache

class StockDataService:
    """股票数据服务类，用于统一管理股票K线数据的获取和处理"""
    
    def __init__(self):
        # 初始化tushare
        self.ts_token = settings.TUSHARE_TOKEN
        ts.set_token(self.ts_token)
        self.pro = ts.pro_api()
        
        # 用于缓存当日数据，避免重复请求
        self._cache: Dict[str, Dict] = {}
        self._cache_time: Dict[str, datetime] = {}
        self.cache_duration = timedelta(minutes=5)  # 缓存有效期5分钟

    def _should_update_cache(self, code: str) -> bool:
        """判断是否需要更新缓存"""
        if code not in self._cache_time:
            return True
        return datetime.now() - self._cache_time[code] > self.cache_duration

    def get_k_data(self, code: str, start_date: Optional[str] = None, 
                   end_date: Optional[str] = None, days: int = 30) -> Optional[pd.DataFrame]:
        """
        获取股票K线数据
        :param code: 股票代码
        :param start_date: 开始日期，格式：YYYYMMDD
        :param end_date: 结束日期，格式：YYYYMMDD
        :param days: 如果未指定开始日期，则获取最近days天的数据
        :return: DataFrame包含OHLCV数据
        """
        try:
            # 检查缓存
            cache_key = f"{code}_{start_date}_{end_date}_{days}"
            if not self._should_update_cache(cache_key):
                return self._cache[cache_key]

            # 处理日期参数
            if end_date is None:
                end_date = datetime.now().strftime('%Y%m%d')
            if start_date is None:
                start_date = (datetime.now() - timedelta(days=days)).strftime('%Y%m%d')

            # 获取数据
            df = self.pro.daily(ts_code=code, start_date=start_date, end_date=end_date)
            
            if df is None or df.empty:
                logging.warning(f"未获取到股票{code}的数据")
                return None

            # 重命名列以匹配原有代码
            df = df.rename(columns={
                'trade_date': 'date',
                'vol': 'volume',
                'ts_code': 'code'
            })
            
            # 按日期升序排序
            df = df.sort_values('date')
            
            # 更新缓存
            self._cache[cache_key] = df
            self._cache_time[cache_key] = datetime.now()

            return df

        except Exception as e:
            logging.error(f"获取股票{code}数据时出错: {str(e)}")
            return None

    def get_real_time_price(self, code: str) -> Optional[float]:
        """
        获取股票实时价格
        :param code: 股票代码
        :return: 最新价格
        """
        try:
            df = ts.get_realtime_quotes(code)
            if df is not None and not df.empty:
                return float(df.iloc[0]['price'])
            return None
        except Exception as e:
            logging.error(f"获取股票{code}实时价格时出错: {str(e)}")
            return None

    def calculate_technical_indicators(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        计算技术指标
        :param df: 包含OHLCV数据的DataFrame
        :return: 添加了技术指标的DataFrame
        """
        try:
            # 计算MACD
            exp1 = df['close'].ewm(span=12, adjust=False).mean()
            exp2 = df['close'].ewm(span=26, adjust=False).mean()
            macd = exp1 - exp2
            signal = macd.ewm(span=9, adjust=False).mean()
            df['macd_hist'] = macd - signal
            df['macd'] = macd
            df['macd_signal'] = signal

            # 计算RSI
            delta = df['close'].diff()
            gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
            loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
            rs = gain / loss
            df['rsi'] = 100 - (100 / (1 + rs))

            # 计算移动平均线
            df['ma5'] = df['close'].rolling(window=5).mean()
            df['ma10'] = df['close'].rolling(window=10).mean()
            df['ma20'] = df['close'].rolling(window=20).mean()

            # 计算成交量均线
            df['volume_ma5'] = df['volume'].rolling(window=5).mean()
            df['volume_ma10'] = df['volume'].rolling(window=10).mean()

            return df

        except Exception as e:
            logging.error(f"计算技术指标时出错: {str(e)}")
            return df 

    def get_stock_basic_info(self, industry: str = None, market_value_min: float = None, 
                            market_value_max: float = None, turnover_rate_min: float = None,
                            pe_min: float = None, pe_max: float = None) -> pd.DataFrame:
        """
        获取符合条件的股票基本信息
        :param industry: 行业
        :param market_value_min: 最小市值（亿）
        :param market_value_max: 最大市值（亿）
        :param turnover_rate_min: 最小换手率
        :param pe_min: 最小市盈率
        :param pe_max: 最大市盈率
        :return: DataFrame包含股票基本信息
        """
        try:
            # 获取股票基本信息
            df = self.pro.stock_basic(exchange='', list_status='L')
            
            # 获取最新交易日的每日指标
            trade_date = self.pro.trade_cal(exchange='', start_date=datetime.now().strftime('%Y%m%d'), 
                                          end_date=datetime.now().strftime('%Y%m%d'), 
                                          is_open='1')['cal_date'].values[0]
            daily_basic = self.pro.daily_basic(trade_date=trade_date)
            
            # 合并基本信息和每日指标
            df = pd.merge(df, daily_basic[['ts_code', 'turnover_rate', 'pe', 'total_mv']], on='ts_code')
            
            # 应用过滤条件
            if industry:
                df = df[df['industry'] == industry]
                
            if market_value_min is not None:
                df = df[df['total_mv'] >= market_value_min * 100000]  # 转换为万元
                
            if market_value_max is not None:
                df = df[df['total_mv'] <= market_value_max * 100000]
                
            if turnover_rate_min is not None:
                df = df[df['turnover_rate'] >= turnover_rate_min]
                
            if pe_min is not None:
                df = df[df['pe'] >= pe_min]
                
            if pe_max is not None:
                df = df[df['pe'] <= pe_max]
            
            return df
            
        except Exception as e:
            logging.error(f"获取股票筛选数据时出错: {str(e)}")
            return None 

    @lru_cache(maxsize=32)
    def get_industry_list(self) -> list:
        """获取所有行业列表"""
        try:
            df = self.pro.stock_basic(exchange='', list_status='L')
            return sorted(df['industry'].unique().tolist())
        except Exception as e:
            logging.error(f"获取行业列表时出错: {str(e)}")
            return [] 